#!/usr/bin/python3
# -*- coding: utf-8 -*-

import os
import time
import numpy as np
import pandas as pd
from glob import glob
from datetime import datetime, timedelta
from joblib import Parallel, delayed
from program.utils  import diff

pd_display_rows  = 20
pd_display_cols  = 8
pd_display_width = 1000
pd.set_option('display.max_rows', pd_display_rows)
pd.set_option('display.min_rows', pd_display_rows)
pd.set_option('display.max_columns', pd_display_cols)
pd.set_option('display.width', pd_display_width)
pd.set_option('display.max_colwidth', pd_display_width)
pd.set_option('display.unicode.ambiguous_as_wide', True)
pd.set_option('display.unicode.east_asian_width', True)
pd.set_option('expand_frame_repr', False)
os.environ['NUMEXPR_MAX_THREADS'] = "256"

from warnings import simplefilter
simplefilter(action='ignore', category=FutureWarning)

from config import pickle_path, data_path, head_columns
from config import factor_class_list
from config import factor_params_list
from config import multi_cal_factors,multi_cal_params
from tqdm import tqdm


def cal_factor_one_param(df,class_name, n, diff_list, _cls):
    '''
    用于计算参数
    '''
    for d_num in diff_list:
        if d_num > 0:
            factor_name = f'{class_name}_bh_{n}_diff_{d_num}'
        else:
            factor_name = f'{class_name}_bh_{n}'
        # 计算因子
        df = df.copy()
        df = _cls.signal(df, n, d_num, factor_name)
    return df[factor_name], factor_name

def cal_one_factor_parallel(df, class_name, params_list, diff_list):
    '''
    并行版本
    '''
    df_cal = df.copy()
    """ ******************** 以下是需要修改的代码 ******************** """
    # =====技术指标
    _cls = __import__('factors.%s' % class_name,  fromlist=('', ))

    n_cal = min(os.cpu_count() - 1 - njobs,len(params_list))
    result = Parallel(n_jobs=n_cal)(delayed(cal_factor_one_param)(df_cal, class_name,n,diff_list,_cls) for n in params_list)
    df_list,factor_list= zip(*result)
    df_list = list(df_list)
    factor_list = list(factor_list)
    df = pd.concat([df]+df_list,axis=1)

    """ ************************************************************ """
    factor_list.sort()
    df = df[head_columns + factor_list]
    df.sort_values(by=['candle_begin_time', ], inplace=True)
    df.reset_index(drop=True, inplace=True)

    return df, class_name


def cal_one_factor(df, class_name, params_list, diff_list):
    df = df.copy()
    """ ******************** 以下是需要修改的代码 ******************** """
    factor_list = []
    # =====技术指标
    _cls = __import__('factors.%s' % class_name,  fromlist=('', ))
    for n in params_list:
        for d_num in diff_list:
            if d_num > 0:
                factor_name = f'{class_name}_bh_{n}_diff_{d_num}'
            else:
                factor_name = f'{class_name}_bh_{n}'
            factor_list.append(factor_name)
            # 计算因子
            df = df.copy()
            # getattr(_cls, 'signal')(df, n, d_num, factor_name)
            df = _cls.signal(df, n, d_num, factor_name)
    """ ************************************************************ """
    factor_list.sort()
    df = df[head_columns + factor_list]
    df.sort_values(by=['candle_begin_time', ], inplace=True)
    df.reset_index(drop=True, inplace=True)

    return df, class_name


def run(trade_type, params_list, diff_list, njobs=16):
    print('\n')
    print(f'trade_type --- {trade_type}')
    # ===创建因子存储目录
    all_factor_path = os.path.join(data_path, trade_type)
    if not os.path.exists(all_factor_path):
        os.makedirs(all_factor_path)

    # ===批量删除头文件
    # 按照下面的处理逻辑, 如果头文件存在就不做覆盖操作
    # 如果重新计算该脚本, 默认头文件不会替换, 可能造成数据错乱
    for header_file in glob(os.path.join(all_factor_path, '*', 'coin_alpha_head.pkl')):
        if os.path.exists(header_file):
            os.remove(header_file)

    # ===开始计算因子
    pbar = tqdm(total=len(glob(os.path.join(pickle_path, f'{trade_type}', '*USDT.pkl'))), position=0, leave=True)
    for pkl_file in glob(os.path.join(pickle_path, f'{trade_type}', '*USDT.pkl')):
        # print('    ', pkl_file)
        symbol = os.path.basename(pkl_file).replace('.pkl', '')
        pbar.set_description(f'计算{symbol}的因子')  # 更新进度条描述信息
        df = pd.read_feather(pkl_file)

        # =====处理原始数据
        df.sort_values(by='candle_begin_time', inplace=True)
        df.drop_duplicates(subset=['candle_begin_time'], inplace=True, keep='last') 
        df['下个周期_avg_price'] = df['avg_price'].shift(-1)  # 计算下根K线开盘买入涨跌幅
        df.reset_index(drop=True, inplace=True)

        if multi_cal_factors:
            if multi_cal_params:
                # 并行计算因子&并行计算参数
                njobs = os.cpu_count() - 1
                results = Parallel(n_jobs=njobs)(delayed(cal_one_factor_parallel)(df, cls_name, params_list, diff_list) for cls_name in factor_class_list)
            else:
                # 并行计算因子&串行计算参数
                njobs = os.cpu_count() - 1
                results = Parallel(n_jobs=njobs)(delayed(cal_one_factor)(df, cls_name, params_list, diff_list) for cls_name in factor_class_list)
        else:
            if multi_cal_params:
                # 串行计算因子&并行计算参数
                results = []
                for cls_name in factor_class_list:
                    # print(cls_name)
                    result = cal_one_factor_parallel(df, cls_name, params_list, diff_list)
                    results.append(result)
            else:
                # 串行计算因子&串行计算参数
                results = []
                for cls_name in factor_class_list:
                    # print(cls_name)
                    result = cal_one_factor(df, cls_name, params_list, diff_list)
                    results.append(result)

        for df, class_name in results:
            symbol_factor_path = os.path.join(all_factor_path, symbol)
            if not os.path.exists(symbol_factor_path):
                os.makedirs(symbol_factor_path)
            # 保存文件头
            head_path = os.path.join(symbol_factor_path, f'coin_alpha_head.pkl')
            if not os.path.exists(head_path):
                df_head = df[head_columns]
                df_head.to_feather(head_path) 
            # 保存因子
            factor_path = os.path.join(symbol_factor_path, f'coin_alpha_factor_{class_name}.pkl')
            df_factors  = df[list(set(df.columns) - set(head_columns))]
            df_factors.to_feather(factor_path) 
        pbar.update(1)  # 更新进度条


if __name__ == '__main__':
    # ===计算因子
    for trade_type in ['spot', 'swap'][1:]:
        print(f'要计算的因子个数:{len(factor_class_list)}')
        factor_num = len(factor_class_list)
        print(f'参数个数:{len(factor_params_list)}')
        njobs = min(factor_num,os.cpu_count() - 2)
        run(trade_type, factor_params_list, diff_list=[0],njobs=njobs)
